from paddleos.img_classification.pre_train import *
from settings import PROJECTS


def train(tf_dict,data_dir,projects_dir,model_str,num_epochs,train_batch_size,lr_decay_epochs,learning_rate):
    train_transforms, eval_transforms = add_transforms(tf_dict)
    train_dataset = set_ImageNet_dataset(data_dir,train_transforms, isshuffle=True)
    eval_dataset = set_ImageNet_dataset(data_dir, eval_transforms)
    model = getattr(pdx.cls, model_str)(num_classes=len(train_dataset.labels))
    model.train(
        num_epochs=10,
        train_dataset=train_dataset,
        train_batch_size=32,
        eval_dataset=eval_dataset,
        lr_decay_epochs=[4, 6, 8],
        learning_rate=0.025,
        save_dir=os.path.join(projects_dir,'output',model_str),
        use_vdl=True)
